package com.m4f.cityclient.oauth.consumer;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;

import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.consumer.OAuth2SecurityContextHolder;
import org.springframework.security.oauth2.consumer.OAuth2SecurityContextImpl;
import org.springframework.web.filter.GenericFilterBean;

public class M4FOAuth2ContextFilter extends GenericFilterBean {

	@Override
	public void doFilter(ServletRequest req, ServletResponse res,
			FilterChain chain) throws IOException, ServletException {
		Authentication authenticationContext = SecurityContextHolder.getContext().getAuthentication();
		// First. Check if the security context is an instance of M4FOAuth2AuthenticationToken. This means
		// that the user has been logged into the application using oauth.
		if(authenticationContext instanceof M4FOAuth2AuthenticationToken) {
			// There is a filter in spring.security.oauth (it's OAuth2ClientContextFilter) that sets the oauth2 context, so here we have to
			// add the access token, used to logged into the application, to that context.
			M4FOAuth2AuthenticationToken auth = (M4FOAuth2AuthenticationToken) authenticationContext;
			OAuth2SecurityContextImpl context = (OAuth2SecurityContextImpl) OAuth2SecurityContextHolder.getContext();
		    if(context != null) { 
		    	HashMap<String, OAuth2AccessToken> tokens = new HashMap<String, OAuth2AccessToken>(context.getAccessTokens());
		    	tokens.put(auth.getProviderAccountId(), auth.getAccessToken());
		    	context.setAccessTokens(Collections.unmodifiableMap(tokens));
			}
		}
		
		chain.doFilter(req, res);
	}

}
